import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import yfinance as yf
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MinMaxScaler
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
import warnings
from datetime import datetime, timedelta

# 设置中文显示
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
warnings.filterwarnings('ignore')


class StockDataCollector:
    """股票数据收集器"""

    def __init__(self, ticker, start_date, end_date=None):
        self.ticker = ticker
        self.start_date = start_date
        self.end_date = end_date if end_date else datetime.now().strftime('%Y-%m-%d')

    def fetch_data(self):
        """从Yahoo Finance获取股票数据"""
        try:
            data = yf.download(self.ticker, start=self.start_date, end=self.end_date)
            print(f"成功获取{self.ticker}从{self.start_date}到{self.end_date}的股票数据")
            return data
        except Exception as e:
            print(f"数据获取失败: {e}")
            return None

    def get_technical_indicators(self, data):
        """计算技术指标"""
        df = data.copy()

        # 移动平均线
        df['MA5'] = df['close'].rolling(window=5).mean()
        df['MA10'] = df['close'].rolling(window=10).mean()
        df['MA20'] = df['close'].rolling(window=20).mean()

        # RSI指标
        delta = df['close'].diff()
        gain = delta.where(delta > 0, 0)
        loss = -delta.where(delta < 0, 0)
        avg_gain = gain.rolling(window=14).mean()
        avg_loss = loss.rolling(window=14).mean()
        rs = avg_gain / avg_loss
        df['RSI'] = 100 - (100 / (1 + rs))

        # MACD指标
        df['EMA12'] = df['close'].ewm(span=12, adjust=False).mean()
        df['EMA26'] = df['close'].ewm(span=26, adjust=False).mean()
        df['MACD'] = df['EMA12'] - df['EMA26']
        df['Signal'] = df['MACD'].ewm(span=9, adjust=False).mean()
        df['Histogram'] = df['MACD'] - df['Signal']

        # 波动率
        df['Volatility'] = df['close'].rolling(window=20).std()

        # 删除包含NaN的行
        df.dropna(inplace=True)
        return df

    def read_csv_to_dataframe(self, file_path):
        """
        从指定的 CSV 文件读取数据到 DataFrame。

        :param file_path: CSV 文件的路径
        :return: 包含 CSV 数据的 DataFrame，如果读取失败则返回 None
        """
        try:
            df = pd.read_csv(file_path)
            print(f"成功读取 {file_path} 的数据")
            return df
        except Exception as e:
            print(f"读取 CSV 文件失败: {e}")
            return None

    @staticmethod
    def select_columns(dataframe, columns):
        """
        从 DataFrame 中取出指定列的数据。

        :param dataframe: 输入的 DataFrame
        :param columns: 要取出的列名，可以是单个列名（字符串）或列名列表
        :return: 包含指定列数据的 DataFrame，如果列名不存在则返回空 DataFrame
        """
        if isinstance(columns, str):
            columns = [columns]
        if all(col in dataframe.columns for col in columns):
            return dataframe[columns]
        else:
            print("部分或全部指定列名不存在于 DataFrame 中。")
            return pd.DataFrame()

class FeatureEngineering:
    """特征工程模块"""

    @staticmethod
    def create_lag_features(df, target_col='close', lags=10, drop_nan=True):
        """创建滞后特征"""
        for i in range(1, lags + 1):
            df[f'{target_col}_lag{i}'] = df[target_col].shift(i)
        if drop_nan:
            df.dropna(inplace=True)
        return df

    @staticmethod
    def create_time_features(df):
        """创建时间特征"""
        df['Day_of_Week'] = df.index.dayofweek
        df['Month'] = df.index.month
        df['Quarter'] = df.index.quarter
        return df

    @staticmethod
    def normalize_features(df, features):
        """归一化特征"""
        scaler = MinMaxScaler()
        df[features] = scaler.fit_transform(df[features])
        return df, scaler


# PyTorch LSTM模型
class LSTMModel(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, output_size, dropout=0.2):
        super(LSTMModel, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers

        self.lstm1 = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, dropout=dropout,
                             bidirectional=False)
        self.dropout = nn.Dropout(dropout)
        self.lstm2 = nn.LSTM(hidden_size, hidden_size, num_layers, batch_first=True, dropout=dropout,
                             bidirectional=False)
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        # 初始化隐藏状态
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device)
        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device)

        # LSTM前向传播
        out, _ = self.lstm1(x, (h0, c0))
        out = self.dropout(out)
        out, _ = self.lstm2(out, (h0, c0))

        # 只取最后一个时间步的输出
        out = self.fc(out[:, -1, :])
        return out


class ModelTrainer:
    """模型训练器"""

    def __init__(self, model_type='random_forest', device='cpu'):
        self.model_type = model_type
        self.model = None
        self.scaler = None
        self.device = device

    def train(self, X, y, test_size=0.2, random_state=42):
        """训练模型"""
        # 划分训练集和测试集
        X_train, X_test, y_train, y_test = train_test_split(
            X, y, test_size=test_size, random_state=random_state, shuffle=False
        )

        # 选择模型
        if self.model_type == 'random_forest':
            self.model = RandomForestRegressor(
                n_estimators=100,
                max_depth=5,
                random_state=random_state
            )
            self.model.fit(X_train, y_train)

        elif self.model_type == 'lstm':
            # 转换为PyTorch张量
            X_train_tensor = torch.FloatTensor(X_train.values).to(self.device)
            y_train_tensor = torch.FloatTensor(y_train.values).to(self.device)
            X_test_tensor = torch.FloatTensor(X_test.values).to(self.device)

            # 重塑数据以适应LSTM [样本, 时间步, 特征]
            X_train_lstm = X_train_tensor.view(X_train_tensor.size(0), X_train_tensor.size(1), 1)
            X_test_lstm = X_test_tensor.view(X_test_tensor.size(0), X_test_tensor.size(1), 1)

            # 创建数据加载器
            train_dataset = TensorDataset(X_train_lstm, y_train_tensor)
            train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

            # 初始化LSTM模型
            input_size = 1
            hidden_size = 50
            num_layers = 1
            output_size = 1

            self.model = LSTMModel(input_size, hidden_size, num_layers, output_size).to(self.device)
            criterion = nn.MSELoss()
            optimizer = optim.Adam(self.model.parameters(), lr=0.001)

            # 训练模型
            self.model.train()
            for epoch in range(50):
                print(f"train epoch: {epoch}")
                for batch_X, batch_y in train_loader:
                    optimizer.zero_grad()
                    outputs = self.model(batch_X)
                    loss = criterion(outputs.squeeze(), batch_y)
                    loss.backward()
                    optimizer.step()

            # 预测
            self.model.eval()
            with torch.no_grad():
                y_pred_tensor = self.model(X_test_lstm)
                y_pred = y_pred_tensor.cpu().numpy().flatten()

        # 模型评估
        y_pred = self.model.predict(X_test) if self.model_type == 'random_forest' else y_pred
        metrics = self.evaluate(y_test, y_pred)

        return {
            'model': self.model,
            'X_train': X_train,
            'X_test': X_test,
            'y_train': y_train,
            'y_test': y_test,
            'y_pred': y_pred,
            'metrics': metrics
        }

    def evaluate(self, y_true, y_pred):
        """评估模型性能"""
        mse = mean_squared_error(y_true, y_pred)
        mae = mean_absolute_error(y_true, y_pred)
        r2 = r2_score(y_true, y_pred)

        return {
            'MSE': mse,
            'MAE': mae,
            'R2 Score': r2,
            'RMSE': np.sqrt(mse)
        }


class Visualizer:
    """可视化模块"""

    @staticmethod
    def plot_prediction(y_true, y_pred, dates, title='股价预测结果'):
        """绘制预测结果"""
        plt.figure(figsize=(12, 6))
        plt.plot(dates, y_true, label='实际价格', color='blue')
        plt.plot(dates, y_pred, label='预测价格', color='red', linestyle='--')
        plt.title(title)
        plt.xlabel('日期')
        plt.ylabel('价格')
        plt.legend()
        plt.grid(True)
        plt.xticks(rotation=45)
        plt.tight_layout()
        plt.show()

    @staticmethod
    def plot_technical_indicators(data, ticker):
        """绘制技术指标"""
        fig, axes = plt.subplots(3, 1, figsize=(14, 12))

        # 绘制收盘价和移动平均线
        axes[0].plot(data.index, data['close'], label='收盘价')
        axes[0].plot(data.index, data['MA5'], label='5日移动平均线')
        axes[0].plot(data.index, data['MA10'], label='10日移动平均线')
        axes[0].plot(data.index, data['MA20'], label='20日移动平均线')
        axes[0].set_title(f'{ticker} 收盘价与移动平均线')
        axes[0].legend()
        axes[0].grid(True)

        # 绘制RSI指标
        axes[1].plot(data.index, data['RSI'])
        axes[1].axhline(70, color='r', linestyle='--')
        axes[1].axhline(30, color='g', linestyle='--')
        axes[1].set_title(f'{ticker} RSI指标')
        axes[1].grid(True)

        # 绘制MACD指标
        axes[2].plot(data.index, data['MACD'], label='MACD')
        axes[2].plot(data.index, data['Signal'], label='Signal')
        axes[2].bar(data.index, data['Histogram'], label='Histogram')
        axes[2].set_title(f'{ticker} MACD指标')
        axes[2].legend()
        axes[2].grid(True)

        plt.tight_layout()
        plt.show()


class StockPredictor:
    """股价预测器 - 整合所有模块"""

    def __init__(self, ticker, start_date, end_date=None, model_type='random_forest', device='cpu'):
        self.ticker = ticker
        self.start_date = start_date
        self.end_date = end_date
        self.model_type = model_type
        self.device = device
        self.data = None
        self.model = None
        self.scaler = None

    def run(self):
        """运行完整的预测流程"""
        # 1. 数据收集
        collector = StockDataCollector(self.ticker, self.start_date, self.end_date)
        self.data = collector.read_csv_to_dataframe(self.ticker)
        if self.data is None:
            return None
        self.data = collector.select_columns(self.data, ['trade_date','open', 'high', 'low', 'close', 'vol'])
        self.data = self.data.sort_values('trade_date', ascending=True)
        self.data = collector.get_technical_indicators(self.data)

        # 2. 特征工程
        fe = FeatureEngineering()
        self.data = fe.create_lag_features(self.data)
        # self.data = fe.create_time_features(self.data)

        # 3. 准备训练数据
        feature_cols = [col for col in self.data.columns if col not in
                        ['trade_date','open', 'high', 'low', 'close', 'vol']]
        target_col = 'close'

        self.data, self.scaler = fe.normalize_features(self.data, feature_cols + [target_col])

        X = self.data[feature_cols]
        y = self.data[target_col]
        print(self.data)

        print("start to train")
        # 4. 模型训练
        trainer = ModelTrainer(self.model_type, self.device)
        result = trainer.train(X, y)
        self.model = result['model']
        print("train finished")

        # 5. 可视化结果
        viz = Visualizer()

        # 绘制技术指标
        viz.plot_technical_indicators(self.data, self.ticker)

        # 绘制预测结果
        test_dates = self.data.index[-len(result['y_test']):]
        viz.plot_prediction(
            result['y_test'],
            result['y_pred'],
            test_dates,
            f'{self.ticker} 股价预测结果 ({self.model_type.upper()})'
        )

        # 6. 预测未来价格
        future_days = 7
        future_predictions = self.predict_future(future_days)

        return {
            'data': self.data,
            'model': self.model,
            'metrics': result['metrics'],
            'future_predictions': future_predictions
        }

    def predict_future(self, days=7):
        """预测未来几天的股价"""
        last_date = self.data['trade_date'].iloc[-1]
        last_date_obj = datetime.strptime(str(last_date), '%Y%m%d')
        future_dates = [last_date_obj + timedelta(days=i+1) for i in range(days)]

        # 获取最后一行数据作为起点
        last_data = self.data.iloc[-1].copy()
        predictions = []

        fe = FeatureEngineering()  # 初始化特征工程类

        for date in future_dates:
            # 准备预测数据
            pred_features = last_data[
                [col for col in self.data.columns if col not in
                 ['open', 'high', 'low', 'close', 'vol', 'trade_date']]
            ].values

            pred_features = pred_features.astype('float64')

            # 预测
            if self.model_type == 'lstm':
                # 转换为PyTorch张量并重塑
                pred_tensor = torch.FloatTensor(pred_features).view(1, len(pred_features), 1).to(self.device)
                self.model.eval()
                with torch.no_grad():
                    pred = self.model(pred_tensor).item()
            else:
                pred = self.model.predict([pred_features])[0]

            print(f"pred: {pred}")
            # 将预测结果添加到列表
            predictions.append(pred)

            # 归一化预测的收盘价
            feature_values = [0] * (len(self.scaler.feature_names_in_) - 1) + [pred]
            feature_values = np.array(feature_values, dtype=np.float32).reshape(1, -1)
            normalized_pred = self.scaler.transform(feature_values)[0][-1]

            # 更新last_data，为下一次预测做准备
            # last_data = last_data.copy()
            last_data['close'] = normalized_pred
            last_data['trade_date'] = date.strftime('%Y%m%d')
            # 注意：这里需要重新应用特征工程，否则滞后特征会丢失
            # 实际应用中可能需要更复杂的逻辑
            self.data = pd.concat([self.data, pd.DataFrame([last_data])])
            self.data = fe.create_lag_features(self.data, drop_nan=False)
            last_data = self.data.iloc[-1].copy()

        # 构造与训练时相同维度的数组进行反归一化
        full_predictions = []
        num_features = len(self.scaler.feature_names_in_)
        for pred in predictions:
            # 初始化全 0 数组
            full_pred = np.zeros(num_features)
            # 找到目标列（close）在训练时的索引
            target_index = np.where(self.scaler.feature_names_in_ == 'close')[0][0]
            full_pred[target_index] = pred
            full_predictions.append(full_pred)

        # 反归一化预测结果
        denorm_full_predictions = self.scaler.inverse_transform(np.array(full_predictions))
        # 提取反归一化后的收盘价
        target_index = np.where(self.scaler.feature_names_in_ == 'close')[0][0]
        denorm_predictions = denorm_full_predictions[:, target_index]
        # pred_features = self.data[-days:].values
        # # 反归一化预测结果
        # denorm_predictions = self.scaler.inverse_transform(
        #     np.array([[p] for p in pred_features]).reshape(-1, 1)
        # ).flatten()

        # 打印未来预测结果
        print("\n未来股价预测:")
        for date, price in zip(future_dates, denorm_predictions):
            print(f"{date.strftime('%Y-%m-%d')}: {price:.2f}")

        return {
            'dates': [date.strftime('%Y-%m-%d') for date in future_dates],
            'prices': denorm_predictions.tolist()
        }


# 使用示例
if __name__ == "__main__":
    # 配置参数
    basepath = "E:\\v\\data\\stock\\"
    ticker = basepath + "newdaily-300866.sz-19900101-20250518.csv"  # 股票代码，如AAPL(苹果), MSFT(微软), BABA(阿里巴巴)
    start_date = "2020-09-01"  # 开始日期
    model_type = "lstm"  # 模型类型: random_forest 或 lstm
    device = "cuda" if torch.cuda.is_available() else "cpu"  # 使用GPU或CPU

    # 创建预测器并运行
    predictor = StockPredictor(ticker, start_date, model_type=model_type, device=device)
    result = predictor.run()

    # 打印评估指标
    if result:
        print("\n模型评估指标:")
        for metric, value in result['metrics'].items():
            print(f"{metric}: {value:.4f}")